package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/go-zookeeper/zk"
)

var (
	zkServers  = flag.String("servers", "localhost:2181", "Zookeeper servers, comma separated")
	operation  = flag.String("op", "", "Operation: create, get, set, delete, exists, ls, watch, tree")
	path       = flag.String("path", "", "Zookeeper node path")
	data       = flag.String("data", "", "Data for create or set operation")
	acl        = flag.String("acl", "open", "ACL: open, creator")
	ephemeral  = flag.Bool("ephemeral", false, "Create ephemeral node")
	sequential = flag.Bool("sequential", false, "Create sequential node")
	version    = flag.Int("version", -1, "Node version for set and delete operations")
	timeout    = flag.Int("timeout", 5, "Connection timeout in seconds")
	retryCount = flag.Int("retry", 3, "Number of connection retries")
	watchTime  = flag.Int("watch-time", 60, "Watch time in seconds for watch operation")
	depthLimit = flag.Int("depth", 10, "Depth limit for tree operation")
	maxNodes   = flag.Int("max-nodes", 1000, "Maximum number of nodes to display in tree operation")
)

func main() {
	flag.Parse()

	// 连接到Zookeeper
	conn, err := connectZookeeper()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to connect to Zookeeper: %v\n", err)
		os.Exit(1)
	}
	defer conn.Close()

	// 执行操作
	switch *operation {
	case "create":
		createNode(conn)
	case "get":
		getNode(conn)
	case "set":
		setNode(conn)
	case "delete":
		deleteNode(conn)
	case "exists":
		existsNode(conn)
	case "ls":
		listChildren(conn)
	case "watch":
		watchNode(conn)
	case "tree":
		listTree(conn)
	default:
		fmt.Fprintf(os.Stderr, "Unknown operation: %s\n", *operation)
		printUsage()
	}
}

func connectZookeeper() (*zk.Conn, error) {
	// 分割服务器地址列表
	servers := strings.Split(*zkServers, ",")

	// 配置连接选项
	conn, eventChan, err := zk.Connect(servers, time.Duration(*timeout)*time.Second)
	if err != nil {
		return nil, err
	}

	// 设置重试计数
	retry := 0
	maxRetries := *retryCount

	// 创建上下文用于超时控制
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(*timeout)*time.Second*(time.Duration(maxRetries)+1))
	defer cancel()

	// 循环等待连接建立或超时
	for {
		select {
		case event := <-eventChan:
			// 打印连接状态变化，便于调试
			fmt.Printf("Connection state changed: %s, path: %s\n", event.State.String(), event.Path)

			// 检查是否已连接
			if event.State == zk.StateConnected {
				fmt.Println("Successfully connected to Zookeeper")
				return conn, nil
			}

			// 检查连接是否已断开
			if event.State == zk.StateDisconnected {
				retry++
				if retry > maxRetries {
					conn.Close()
					return nil, fmt.Errorf("connection failed after %d retries", maxRetries)
				}
				fmt.Printf("Connection disconnected, retrying (%d/%d)...\n", retry, maxRetries)
				time.Sleep(time.Second * 2) // 等待一段时间后重试
			}

			// 检查是否认证失败
			if event.State == zk.StateAuthFailed {
				conn.Close()
				return nil, fmt.Errorf("authentication failed")
			}

		case <-ctx.Done():
			conn.Close()
			return nil, fmt.Errorf("connection timeout after %d seconds", *timeout*(maxRetries+1))
		}
	}
}

func getACL() []zk.ACL {
	if *acl == "creator" {
		return zk.DigestACL(zk.PermAll, "user", "password")
	}
	// 默认使用开放ACL
	return zk.WorldACL(zk.PermAll)
}

func getCreateFlags() int32 {
	var flags int32
	if *ephemeral {
		flags |= zk.FlagEphemeral
	}
	if *sequential {
		flags |= zk.FlagSequence
	}
	return flags
}

func createNode(conn *zk.Conn) {
	if *path == "" {
		fmt.Fprintln(os.Stderr, "Path is required for create operation")
		os.Exit(1)
	}

	flags := getCreateFlags()
	aclList := getACL()

	createdPath, err := conn.Create(*path, []byte(*data), flags, aclList)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create node: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Node created successfully: %s\n", createdPath)
}

func getNode(conn *zk.Conn) {
	if *path == "" {
		fmt.Fprintln(os.Stderr, "Path is required for get operation")
		os.Exit(1)
	}

	data, stat, err := conn.Get(*path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to get node: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Data: %s\n", string(data))
	fmt.Printf("Stat: \n")
	fmt.Printf("  Version: %d\n", stat.Version)
	fmt.Printf("  Aversion: %d\n", stat.Aversion)
	fmt.Printf("  Cversion: %d\n", stat.Cversion)
	fmt.Printf("  Ctime: %s\n", time.Unix(stat.Ctime/1000, 0))
	fmt.Printf("  Mtime: %s\n", time.Unix(stat.Mtime/1000, 0))
	fmt.Printf("  DataLength: %d\n", stat.DataLength)
	fmt.Printf("  NumChildren: %d\n", stat.NumChildren)
}

func setNode(conn *zk.Conn) {
	if *path == "" {
		fmt.Fprintln(os.Stderr, "Path is required for set operation")
		os.Exit(1)
	}

	stat, err := conn.Set(*path, []byte(*data), int32(*version)) // 使用指定版本
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to set node: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Node updated successfully\n")
	fmt.Printf("New version: %d\n", stat.Version)
}

func deleteNode(conn *zk.Conn) {
	if *path == "" {
		fmt.Fprintln(os.Stderr, "Path is required for delete operation")
		os.Exit(1)
	}

	err := conn.Delete(*path, int32(*version)) // 使用指定版本
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to delete node: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Node deleted successfully: %s\n", *path)
}

func existsNode(conn *zk.Conn) {
	if *path == "" {
		fmt.Fprintln(os.Stderr, "Path is required for exists operation")
		os.Exit(1)
	}

	exists, stat, err := conn.Exists(*path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to check node existence: %v\n", err)
		os.Exit(1)
	}

	if exists {
		fmt.Printf("Node exists: %s\n", *path)
		fmt.Printf("Stat: \n")
		fmt.Printf("  Version: %d\n", stat.Version)
		fmt.Printf("  DataLength: %d\n", stat.DataLength)
		fmt.Printf("  NumChildren: %d\n", stat.NumChildren)
	} else {
		fmt.Printf("Node does not exist: %s\n", *path)
	}
}

func listChildren(conn *zk.Conn) {
	if *path == "" {
		fmt.Fprintln(os.Stderr, "Path is required for ls operation")
		os.Exit(1)
	}

	children, stat, err := conn.Children(*path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to list children: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Children of %s (total: %d):\n", *path, len(children))
	for _, child := range children {
		fmt.Printf("  %s\n", child)
	}
	fmt.Printf("Stat: \n")
	fmt.Printf("  NumChildren: %d\n", stat.NumChildren)
}

func watchNode(conn *zk.Conn) {
	if *path == "" {
		fmt.Fprintln(os.Stderr, "Path is required for watch operation")
		os.Exit(1)
	}

	fmt.Printf("Watching node %s for %d seconds...\n", *path, *watchTime)

	watchCtx, watchCancel := context.WithTimeout(context.Background(), time.Duration(*watchTime)*time.Second)
	defer watchCancel()

	// 创建一个通道接收watch事件
	// Removed unused eventChan declaration

	// 设置监听器函数
	// Removed unused watchFunc declaration

	// 开始监听节点
	exists, _, eventChan, err := conn.ExistsW(*path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to start watching node: %v\n", err)
		os.Exit(1)
	}

	if exists {
		fmt.Printf("Node %s exists, waiting for changes...\n", *path)
	} else {
		fmt.Printf("Node %s does not exist, waiting for creation...\n", *path)
	}

	// 设置事件回调
	// Directly handle the watch event using the ExistsW method

	// 等待事件或超时
	select {
	case event := <-eventChan:
		fmt.Printf("Watch event received:\n")
		fmt.Printf("  Type: %s\n", event.Type.String())
		fmt.Printf("  State: %s\n", event.State.String())
		fmt.Printf("  Path: %s\n", event.Path)
		fmt.Printf("  Err: %v\n", event.Err)
	case <-watchCtx.Done():
		fmt.Printf("Watch timeout after %d seconds\n", *watchTime)
	}
}

func printUsage() {
	fmt.Fprintln(os.Stderr, "Usage: zk-tool -servers=localhost:2181 -op=create -path=/test -data=testdata")
	fmt.Fprintln(os.Stderr, "\nAvailable operations:")
	fmt.Fprintln(os.Stderr, "  create  - Create a new node")
	fmt.Fprintln(os.Stderr, "  get     - Get node data and stat")
	fmt.Fprintln(os.Stderr, "  set     - Set node data")
	fmt.Fprintln(os.Stderr, "  delete  - Delete a node")
	fmt.Fprintln(os.Stderr, "  exists  - Check if a node exists")
	fmt.Fprintln(os.Stderr, "  ls      - List children of a node")
	fmt.Fprintln(os.Stderr, "  watch   - Watch node for changes")
	fmt.Fprintln(os.Stderr, "  tree    - List directory tree recursively")
	fmt.Fprintln(os.Stderr, "\nOptions:")
	flag.PrintDefaults()
}

// listTree 递归列出目录树信息
func listTree(conn *zk.Conn) {
	if *path == "" {
		fmt.Fprintln(os.Stderr, "Path is required for tree operation")
		os.Exit(1)
	}

	// 检查节点是否存在
	exists, _, err := conn.Exists(*path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to check node existence: %v\n", err)
		os.Exit(1)
	}

	if !exists {
		fmt.Printf("Node does not exist: %s\n", *path)
		existsNode(conn)
		return
	}

	fmt.Printf("Directory tree for %s (depth limit: %d, max nodes: %d):\n", *path, *depthLimit, *maxNodes)

	// 使用闭包来跟踪节点计数
	nodeCount := 0
	maxNodesReached := false

	// 递归函数来列出树
	var recursiveList func(path string, prefix string, depth int)
	recursiveList = func(path string, prefix string, depth int) {
		// 检查是否达到最大节点数
		nodeCount++
		if nodeCount > *maxNodes {
			maxNodesReached = true
			return
		}

		// 获取节点状态
		exists, stat, err := conn.Exists(path)
		if !exists {
			fmt.Printf("%s[INFO] Node does not exist: %s\n", prefix, path)
			return
		}
		if err != nil {
			fmt.Printf("%s[ERROR] %s: %v\n", prefix, path, err)
			return
		}

		// 打印当前节点
		nodeType := "[DIR]"
		if stat.NumChildren == 0 {
			nodeType = "[FILE]"
		}
		fmt.Printf("%s%s %s (ver: %d, data: %d bytes)\n", prefix, nodeType, path, stat.Version, stat.DataLength)

		// 检查是否达到深度限制
		if depth >= *depthLimit {
			return
		}

		// 获取子节点
		children, _, err := conn.Children(path)
		if err != nil {
			fmt.Printf("%s  [ERROR] Failed to list children: %v\n", prefix, err)
			return
		}

		// 递归列出子节点
		for i, child := range children {
			// 检查是否达到最大节点数
			if maxNodesReached {
				return
			}

			childPath := path
			if childPath != "/" {
				childPath += "/"
			}
			childPath += child

			// 确定前缀
			childPrefix := prefix
			if i == len(children)-1 {
				childPrefix += "└── "
			} else {
				childPrefix += "├── "
			}

			// 递归调用
			recursiveList(childPath, childPrefix, depth+1)
		}
	}

	// 开始递归列出
	recursiveList(*path, "", 0)

	// 如果达到了最大节点数，提示用户
	if maxNodesReached {
		fmt.Printf("\nMax nodes limit (%d) reached. Use -max-nodes flag to display more nodes.\n", *maxNodes)
	}

	fmt.Printf("\nTotal nodes displayed: %d\n", nodeCount)
}
